//
// Created by 郝羽 on 2017/12/11.
//

#ifndef RTC_XSWITCH2_AIMD_RATE_CONTROL_HPP
#define RTC_XSWITCH2_AIMD_RATE_CONTROL_HPP

#include <algorithm>
#include <cmath>
#include <cstdint>

#include "bwe/bwe.hpp"

namespace bwe {


    enum class RateControlState {
        Hold,
        Increase,
        Decrease,
    };

    enum class RateControlRegion {
        NearMax,
        AboveMax,
        Unknown,
    };

    class AimdRateControl {
    public:
        bool HasBps() const {
            return has_bps_;
        };

        // for unit test
        void SetEstimate(int bps, int64_t now_ms) {
            has_bps_ = true;
            cur_bps_ = clamp(bps, bps);
            last_change_ms_ = now_ms;
        };

        // for unit test
        void SetStartBitrate(int bps) {
            has_bps_ = true;
            cur_bps_ = bps;
        };

        // for unit test
        int GetNearMaxIncreaseRateBps(int rtt_ms) const {
            double bits_per_frame = static_cast<double>(cur_bps_) / 30.0;
            double packets_per_frame = std::ceil(bits_per_frame / (8.0 * 1200.0));
            double avg_packet_size_bits = bits_per_frame / packets_per_frame;

            // Approximate the over-use estimator delay to 100 ms.
            const int64_t response_time = (rtt_ms + 100) * 2;
            constexpr double kMinIncreaseRateBps = 4000;
            return static_cast<int>(std::max(kMinIncreaseRateBps, (avg_packet_size_bits * 1000) / response_time));
        };

        // for unit test
        int GetExpectedBandwidthPeriodMs() const {
            const int kMinPeriodMs = 2000;
            constexpr int kDefaultPeriodMs = 3000;
            constexpr int kMaxPeriodMs = 50000;

            int increase_rate = GetNearMaxIncreaseRateBps(200);
            if (!last_decrease_)
                return kDefaultPeriodMs;

            return std::min(kMaxPeriodMs, std::max<int>(1000 * static_cast<int64_t>(last_decrease_) / increase_rate, kMinPeriodMs));
        }

        void SetMinBitrate(uint32_t min) {
            min_bps_ = min;
        };

        bool ReduceFurther(int64_t now_ms, uint32_t incoming_bps) const {
            const int64_t reduction_interval_ms = 100;
            if (now_ms - last_change_ms_ >= reduction_interval_ms) {
                return true;
            }
            if (HasBps()) {
                const uint32_t threshold = static_cast<uint32_t>(0.5 * LatestEstimate());
                return incoming_bps < threshold;
            }
            return false;
        };

        uint32_t LatestEstimate() const {
            return cur_bps_;
        };

        uint32_t Update(const RateControlInput& input, int64_t now_ms) {
            if (!HasBps()) {
                if (first_update_ms_ < 0) {
                    if (input.has_bitrate) {
                        first_update_ms_ = now_ms;
                    }
                } else if (now_ms - first_update_ms_ > 5000 && input.has_bitrate) {
                    has_bps_ = true;
                }
            }
            cur_bps_ = changeBitrate(cur_bps_, input, now_ms);
            return cur_bps_;
        };

    private:
        uint32_t changeBitrate(uint32_t new_bps, const RateControlInput& input, int64_t now_ms) {
            uint32_t incoming_bps = cur_bps_;
            if (input.has_bitrate) {
                incoming_bps = input.bitrate;
            }
            if (!HasBps() && input.usage != Usage::OverUsing) {
                return cur_bps_;
            }
            changeState(input, now_ms);
            const float incoming_kbps = incoming_bps / 1000.0f;
            const float std_max = sqrt(var_max_kbps_ * avg_max_kbps_);
            switch (state_) {
                case RateControlState::Hold:
                    // patch over webrtc
                    if (incoming_bps > cur_bps_) {
                        cur_bps_ += (incoming_bps - cur_bps_) * 0.2;
                    }
                    break;
                case RateControlState::Increase:
                    if (avg_max_kbps_ >=0 && incoming_kbps > avg_max_kbps_ + 3 * std_max) {
                        changeRegion(RateControlRegion::Unknown);
                        avg_max_kbps_ = -1.0f;
                    }
                    if (region_ == RateControlRegion::NearMax) {
                        uint32_t add_bps = additiveIncrease(now_ms);
                        new_bps += add_bps;

                    } else {
                        uint32_t add_bps = multiplicativeIncrease(now_ms);
                        new_bps += add_bps;
                    }
                    last_change_ms_ = now_ms;
                    break;
                case RateControlState::Decrease:
                    new_bps = static_cast<uint32_t>(0.85 * incoming_bps + 0.5);
                    if (new_bps > cur_bps_) {
                        if (region_ != RateControlRegion::Unknown) {
                            new_bps = static_cast<uint32_t>(0.85 * avg_max_kbps_ * 1000 + 0.5f);
                        }
                        new_bps = std::min(new_bps, cur_bps_);
                    }

                    changeRegion(RateControlRegion::NearMax);
                    if (HasBps() && incoming_bps < cur_bps_) {
                        last_decrease_ = cur_bps_ - new_bps;
                    }
                    if (incoming_kbps < avg_max_kbps_ - 3 * std_max) {
                        avg_max_kbps_ = -1.0f;
                    }
                    has_bps_ = true;
                    updateMaxBitrate(incoming_kbps);
                    state_ = RateControlState::Hold;
                    last_change_ms_ = now_ms;
                    break;
            }
            return clamp(new_bps, incoming_bps);
        };


        void changeState(const RateControlInput& input, int64_t now_ms) {
            switch (input.usage) {
                case Usage::Normal:
                    if (state_ == RateControlState::Hold) {
                        last_change_ms_ = now_ms;
                        state_ = RateControlState::Increase;
//                        std::cout << "|normal| hold -> increase" << std::endl;
                    }
                    break;
                case Usage::OverUsing:
                    if (state_ != RateControlState::Decrease) {
//                        std::cout << "|overusing| " << (state_ == RateControlState::Hold ? "hold" : "increase") << " -> decrease" << std::endl;
                        state_ = RateControlState::Decrease;
                    }
                    break;
                case Usage::UnderUsing:
//                    std::cout << "|underusing| " << (state_ == RateControlState::Hold ? "hold" : (state_ == RateControlState::Increase ? "increase" : "decrease")) << " -> hold" << std::endl;
                    state_ = RateControlState::Hold;

                    break;
            }
        };

        void changeRegion(const RateControlRegion& region) {
            region_ = region;
        };

        void updateMaxBitrate(float kbps) {
            const float alpha = 0.05f;
            if (avg_max_kbps_ == -1.0f) {
                avg_max_kbps_ = kbps;
            } else {
                avg_max_kbps_ = (1 - alpha) * avg_max_kbps_ + alpha * kbps;
            }
            // const float norm = std::max(avg_max_kbps_, 1.0f);
            var_max_kbps_ = (1 - alpha) * var_max_kbps_ + alpha * (avg_max_kbps_ - kbps) * (avg_max_kbps_ - kbps);
            if (var_max_kbps_ < 0.4f) {
                var_max_kbps_ = 0.4f;
            } else if (var_max_kbps_ > 2.5f) {
                var_max_kbps_ = 2.5f;
            }
        };

        uint32_t clamp(uint32_t new_bps, uint32_t incoming_bps) const {
            const uint32_t max_bps = static_cast<uint32_t>(1.5f * incoming_bps) + 10000;
            if (new_bps > cur_bps_ && new_bps > max_bps) {
                new_bps = std::max(cur_bps_, max_bps);
            }
            new_bps = std::max(new_bps, min_bps_);
            return new_bps;
        };

        uint32_t additiveIncrease(int64_t now_ms) {
            return static_cast<uint32_t>((now_ms - last_change_ms_) * nearMaxIncreaseBps() / 1000);
        };

        uint32_t nearMaxIncreaseBps() {
            double bits_per_frame = static_cast<double>(cur_bps_) / 30.0;
            double packets_per_frame = std::ceil(bits_per_frame / (8.0 * 1200.0));
            double avg_packet_size_bits = bits_per_frame / packets_per_frame;

            const int64_t rtt = 100;
            const double minIncreaseBps = 4000;
            return static_cast<uint32_t>(std::max(minIncreaseBps, (avg_packet_size_bits * 1000) / (rtt + 100)));
        };

        uint32_t multiplicativeIncrease(int64_t now_ms) {
            double alpha = 1.08;
            if (last_change_ms_ > -1) {
                int64_t time_since_last_change = std::min(now_ms - last_change_ms_, (int64_t)1000);
                alpha = std::pow(alpha, time_since_last_change / 1000.0);
            }
            uint32_t delta = std::max(static_cast<uint32_t>(cur_bps_ * (alpha - 1) + 0.5), 1000u);
            return delta;
        };

    private:
        bool has_bps_ = false;
        uint32_t cur_bps_ = 300000;
        uint32_t min_bps_ = 0;
        int64_t first_update_ms_ = -1;
        int64_t last_change_ms_ = -1;
        RateControlState state_ = RateControlState::Hold;
        RateControlRegion region_ = RateControlRegion::Unknown;
        float avg_max_kbps_ = -1.0f;
        float var_max_kbps_ = -0.4f;
        uint32_t last_decrease_ = 0;
    };
}

#endif //RTC_XSWITCH2_AIMD_RATE_CONTROL_HPP
